!// Dmitry.Konovalov@jcu.edu.au Dmitry.A.Konovalov@gmail.com, 2013

!! 19-Mar-2013: trying the syle from:
!! http://fortran90.org/src/best-practices.html
!!Only use modules and programs. Always setup a module in the following way:
module lgrr_mod;
use func_mod;  
use qdr_mod;  
use gamma_mod;
implicit none;
logical, save :: LGRR_TRACE = .TRUE.;

private;

! NOTE 9-Apr-2013: name any new public functions as lgrr_xxx()
public test_lgrr_mod, lgrr_arr, lgrr_lcr, lgrr_orth, LGRR_OPT_T &
&, lgrr_lcr3, lgrr_lcr_L;


type LGRR_OPT_T 
 real(DP) :: alpha, lambda;
 integer :: maxn, L;
end type


contains;
subroutine test_lgrr_mod();  
  character(LEN=64) :: name = "test_lgrr_mod";         
  call dbg_in(name);
  call test_lgrr();
  call test_lgrr_orth();
  call test_ho3d_orth();!Three-dimensional harmonic-oscillator (HO3D) wave functions
  call test_lgrr_lcr();
  call dbg_out(name); 
end subroutine


subroutine test_lgrr();   
  real(DP), pointer :: xgrid(:), av1(:,:), av13(:,:);     ! av - array of vectors  
  real(DP) :: alpha;
  integer n, nx;
  character(LEN=64) :: name = "test_lgrr";         
  call dbg_in(name);
  
  alpha = D0;
  nx = 11;
  xgrid => mem_vec(nx);
  call grid_step(xgrid,  D0, D1); call vec_dbg("xgrid", xgrid);

  av13 => mem_veca(xgrid, 13);
  call lgrr_arr(av13, xgrid, alpha);  call mtrx_dbg("lgrr_arr(xgrid, av, opt)=", av13);
  
  av1 => mem_veca(xgrid, 1);
  call lgrr_arr(av1, xgrid, alpha);  call mtrx_dbg("lgrr_arr(xgrid, av1, opt)=", av1);
  
  n = 0;
  call dbg_assert_abs_diff_LE("L_0(0)=", av1(1, n+1), D1, EPS32);
  call dbg_assert_abs_diff_LE("L_0(0)=", av13(1, n+1), D1, EPS32);
  call dbg_assert_abs_diff_LE("L_0(1)=", av1(nx, n+1), D1, EPS32);
  call dbg_assert_abs_diff_LE("L_0(1)=", av13(nx, n+1), D1, EPS32);
  
  n = 1;
  call dbg_assert_abs_diff_LE("L_1(0)=", av13(1, n+1), D1, EPS32);
  call dbg_assert_abs_diff_LE("L_1(0.5)=", av13(6, n+1), D05, EPS32);
  call dbg_assert_abs_diff_LE("L_1(1)=", av13(nx, n+1), D0, EPS32);
  
  n = 2;
  call dbg_assert_abs_diff_LE("L_2(1)=", av13(nx, n+1), -D05, EPS32);
  
  n = 3;
  call dbg_assert_abs_diff_LE("L_3(1)=", av13(nx, n+1), -D2/D3, EPS32);

  n = 4;
  call dbg_assert_abs_diff_LE("L_4(1)=", av13(nx, n+1), -0.625_DP, EPS32);
  
  n = 5;
  call dbg_assert_abs_diff_LE("L_5(1)=", av13(nx, n+1), -0.4666666667_DP, EPS10);
  call dbg_assert_abs_diff_LE("L_5(0.5)=", av13(6, n+1), -0.4455729167_DP, EPS10);

  n = 9;
  call dbg_assert_abs_diff_LE("L_9(1)=", av13(nx, n+1), 0.3097442681_DP, EPS10);
  
  n = 10;
  call dbg_assert_abs_diff_LE("L_10(1)=", av13(nx, n+1), 0.4189459325_DP, EPS10);
  
  n = 12;
  call dbg_assert_abs_diff_LE("L_12(1)=", av13(nx, n+1), 0.4962122235_DP, EPS11);
  
  call mem_del_vec(xgrid);
  call mem_del_mtrx2(av1, av13);
  call dbg_out(name); 
end subroutine

subroutine test_lgrr_orth();   
  real(DP), pointer :: xgrid(:), xqdr(:), av(:,:);     ! av - array of vectors  
  type(LGRR_OPT_T) :: opt;
  integer :: nx;
  real(DP) :: max_err, norm; 
  character(LEN=64) :: name = "test_lgrr_orth";         
  call dbg_in(name);

  opt%alpha = D2;  opt%lambda = D07;  opt%maxn = 4;
  nx = 501;
  xgrid => mem_vec(nx);
  call grid_step(xgrid,  D0, D60); call vec_dbg("xgrid", xgrid);
  xqdr => mem_vec(nx);  
  call qdr_step_p5(xqdr, xgrid); call vec_dbg("xqdr", xqdr); ! x-quadratures
 
  ! LGRR_ORTH
  call mem_del_mtrx(av);
  av => mem_veca(xgrid, opt%maxn + 1);
  call lgrr_orth(av, xgrid, opt);  call mtrx_dbg("lgrr_orth(xgrid, av, opt)=", av);
  norm = qdr_intgrl2(xqdr, av(:,1), av(:,1));
  call dbg_assert_abs_diff_LE("norm orth L_0=", norm, D1, EPS8);
  
  max_err = qdr_max_orth_err(xqdr, av);
  call dbg_assert_abs_diff_LE("qdr_max_orth_err=", max_err, D0, EPS5); !NOTE: max err = E-5
 
  call mem_del_vec2(xgrid, xqdr);
  call mem_del_mtrx(av);
  call dbg_out(name); 
end subroutine

!Three-dimensional harmonic-oscillator (HO3D) wave functions
subroutine test_ho3d_orth();   
  real(DP), pointer :: rv(:), rqdr(:), av(:,:);     ! av - array of vectors  
  type(LGRR_OPT_T) :: opt;
  type(GRID_LCR_T) :: lcr;
  real(DP) :: max_err, norm; 
  character(LEN=64) :: name = "test_ho3d_orth";         
  call dbg_in(name);
  
  call grid_lcr_init(lcr);
  call dbg_dp("lcr%rmin", lcr%rmin);

  lcr%xmin = -D5;  lcr%rmax = D20;  lcr%nx = 501;
  call grid_lcr(lcr);  !call vec_dbg("lcr", xgrid);
  call qdr_lcr(lcr);
  
  opt%alpha = D2;  opt%lambda = D07;  opt%maxn = 4;
  rv => mem_vec(lcr%nx);
  call grid_step(rv, lcr%rmin, lcr%rmax); call vec_dbg("rv", rv);
  rqdr => mem_vec_from(rv);  
  call qdr_step_p5(rqdr, rv); call vec_dbg("rqdr", rqdr); ! x-quadratures
 
 ! ho3d_lcr
  av => mem_veca(lcr%xv, opt%maxn + 1);
  call ho3d_lcr(av, lcr, opt);  call mtrx_dbg("ho3d_lcr(rv, av, opt)=", av);
  norm = qdr_intgrl_lcr2(lcr, av(:,1), av(:,1));
  call dbg_assert_abs_diff_LE("norm ho3d_lcr L_0=", norm, D1, EPS14); !was EPS8 with r-grid
  max_err = qdr_max_orth_err(lcr%xqdr_cr2, av);
  call dbg_assert_abs_diff_LE("ho3d_lcr max_orth_err=", max_err, D0, EPS12); !was = EPS8
 
  ! ho3d_orth
  call mem_del_mtrx(av);
  av => mem_veca(rv, opt%maxn + 1);
  call ho3d_orth(av, rv, opt);  call mtrx_dbg("ho3d_orth(xgrid, av, opt)=", av);
  norm = qdr_intgrl2(rqdr, av(:,1), av(:,1));
  call dbg_assert_abs_diff_LE("norm ho3d L_0=", norm, D1, EPS8);
  
  max_err = qdr_max_orth_err(rqdr, av);
  call dbg_assert_abs_diff_LE("qdr_max_orth_err=", max_err, D0, EPS8); !NOTE: max err = E-8
 
  call mem_del_vec2(rv, rqdr);
  call mem_del_mtrx(av);
  call grid_lcr_init(lcr);
  call dbg_out(name); 
end subroutine

subroutine test_lgrr_lcr();   
  real(DP), pointer :: av(:,:);     ! av - array of vectors  
  type(LGRR_OPT_T) :: opt;
  type(GRID_LCR_T) :: lcr;
  real(DP) :: max_err, norm; 
  character(LEN=64) :: name = "test_lgrr_lcr";         
  call dbg_in(name);

  opt%alpha = D2;  opt%lambda = D07;  opt%maxn = 4;
  call grid_lcr_init(lcr);
  lcr%xmin = -D5;  lcr%rmax = D100;  lcr%nx = 501;
  !!! NOTE: using 100. as rmax
  call grid_lcr(lcr);  !call vec_dbg("lcr", xgrid);
  call qdr_lcr(lcr);
  
  ! LGRR_LCR
  call mem_del_mtrx(av);
  av => mem_veca(lcr%xv, opt%maxn + 1);
  call lgrr_lcr(av, lcr, opt);  call mtrx_dbg("lgrr_arr(xgrid, av, opt)=", av);
  norm = qdr_intgrl_lcr2(lcr, av(:,1), av(:,1));
  call dbg_assert_abs_diff_LE("test_lgrr_lcr, norm orth L_0=", norm, D1, EPS13); !was EPS8 with r-grid
  
  max_err = qdr_max_orth_err(lcr%xqdr_cr2, av);
  call dbg_assert_abs_diff_LE("qdr_max_orth_err=", max_err, D0, EPS12); !was = EPS5
 
  call mem_del_mtrx(av);
  call grid_lcr_init(lcr);
  call dbg_out(name); 
end subroutine

SUBROUTINE lgrr_n1(res, x, alpha);     
  real(DP), intent(OUT) :: res(:);      ! vector-x, vector-y
  real(DP), intent(IN)  :: x(:), alpha;   
  call dbg_str("-->lgrr_n1()", LGRR_TRACE);
  call vec_init_res(res, x);
  res = (alpha + D1);
  res = res - x;
end SUBROUTINE


! NOTE: this is for lambda=1; call with xgrid*lambda 
subroutine lgrr_arr(av, xgrid, alpha);   
  real(DP), intent(OUT) :: av(:,:);     ! av - array of vectors  
  real(DP), intent(IN) :: xgrid(:), alpha;      
  real(DP) :: x, Ln, Ln1, Ln2, dn;
  integer :: n, ix, minx, maxx, maxn, nx;
  integer :: minr, maxr, nr;  ! rows
  integer :: minc, maxc, nc;  ! cols
  call dbg_str("-->lgrr_arr()", LGRR_TRACE); 
  call vec_range(minx, maxx, xgrid, nx);
  call mtrx_i1_range(minr, maxr, av, nr);
  call mtrx_i2_range(minc, maxc, av, nc);
  call dbg_assert_true("nr > 0", nr > 0);
  call dbg_assert_true("nc > 0", nc > 0);
  call dbg_assert_true("nx > 0", nx > 0);
  call dbg_assert_true("minr == minx", minr == minx);
  call dbg_assert_true("maxr == maxx", maxr == maxx);

  maxn = nc - 1;  !note: n=0,1,...,

  av(:, 1) = D1;  ! n=0
  if (maxn == 0)  return;

  call lgrr_n1(av(:, 2), xgrid, alpha);  ! n = 1
  if (maxn == 1)  return;
  
  do ix = minx, maxx
    x = xgrid(ix);
    Ln2 = av(ix, 1); ! L_{n-2}
    Ln1 = av(ix, 2); ! L_{n-1}
    do n = 2, maxn   
      dn = real(n, DP);  ! d-real version of n
      !! (n+1) L_(n+1) = ((2n+a+1)-x) L_n   - (n+a)   L_n_1  // from p588 of russian Abramowitz
      !!     n L_n     = ((2n+a-1)-x) L_n_1 - (n-1+a) L_n_2    end do
      Ln = ((D2 * dn + alpha - D1) - x) * Ln1 - (dn - D1 + alpha) * Ln2; 
      Ln = Ln / dn;
      av(ix, n+1) = Ln;
      Ln2 = Ln1;
      Ln1 = Ln;
    end do
  end do  
end subroutine


!//  Orthonormal version: pp579-580 from Russian copy of Abramowitz
!//                               (a)
!//   LagrrBasis's polinomials   L  (y), n = 0,..., nmax
!//                               n
!//   oo   (a)    (a)    -y  a
!//   I   L  (y) L  (y) e   y  dy = Gamma(a + n + 1) / n!
!//   0     n     n'
!//
!//     y = lambda * r; dr = dy/lambda
!//
!//   oo   (a)             (a)             -y  a
!//   I   L  (lambda * r) L  (lambda * r) e   y  dr = 1/lambda * Gamma(a + n + 1) / n!
!//   0     n               n'
!//                                                oo
!//Set of orthonormal functions L_n(r), such that I L_n(r) L_n'(r) dr = delta_nn'
!//                                                0
!// NOTE: integration is dr  not  r**2 dr
!//
!// P(r) = NORM x^{a/2} exp(-x/2) L_n^{a}(x);    x=lambda*r
!// NORM = sqrt(lambda) * sqrt(GAMMA(p+1) / GAMMA(p+a+1))
!// a = 2L+2;
subroutine lgrr_orth(av, xgrid, opt);  
  real(DP), intent(OUT) :: av(:,:);     ! av - array of vectors  
  real(DP), intent(IN) :: xgrid(:);      
  type(LGRR_OPT_T), intent(IN) :: opt;
  integer :: i, mini, maxi;
  real(DP) :: norm;
  character(LEN=64) :: name = "lgrr_orth";         
  call dbg_in(name);
  call lgrr_orth_u1(av, opt%lambda * xgrid, opt%alpha);  
  call mtrx_i2_range(mini, maxi, av);
  norm = SQRT(opt%lambda); 
  do i = mini, maxi  
    av(:, i) = av(:, i) * norm; 
  end do 
end subroutine

! n=0; NOTE: this is for lambda=1; use lambda*xv for others
subroutine lgrr_orth_u1(av, xgrid, alpha);  
  real(DP), intent(OUT) :: av(:,:);     ! av - array of vectors  
  real(DP), intent(IN) :: xgrid(:), alpha;      
  integer :: i, mini, maxi, n;
  real(DP) :: norm;
  character(LEN=64) :: name = "lgrr_orth_u1";         
  call dbg_in(name);
  
  call lgrr_arr(av, xgrid, alpha);
  call mtrx_i2_range(mini, maxi, av);
  
  ! NOTE!!! av(:, 1) is just 1. for n=0
  ! so use this as the working space
  call lgrr_orth_n0(av(:, 1), xgrid, alpha);  
  !call vec_dbg("lgrr_orth_n0(av(:, 1)", av(:, 1), LGRR_TRACE);
  do i = 2, maxi  ! NOTE 2
    av(:, i) = av(:, i) * av(:, 1);
    !call vec_dbg("av(:, i) * av(:, 1)", av(:, i), LGRR_TRACE);
  end do 
  do i = 1, maxi  
    n = i - 1;
    norm = lgrr_orth_norm(n, alpha); !call dbg_dp("norm", norm, LGRR_TRACE);
    !call vec_dbg("av(:, i)", av(:, i), LGRR_TRACE);
    av(:, i) = av(:, i) * norm; 
    !call vec_dbg("av(:, i)*norm", av(:, i), LGRR_TRACE);
  end do 
end subroutine


!!// NOTE: integration is dr  not  r**2 dr
!!// P(r) = NORM x^{L+1} exp(-x^2/2) L_n^{L+1/2}(x^2);    x=(lambda*r)
!!// P(r) = NORM y^{a/2} x^{1/2} exp(-y/2) L_n^{a}(y);    y=x^2; a=L+1/2
!!// NORM = sqrt(2) sqrt( GAMMA(p+1) / GAMMA(p+a+1))
subroutine ho3d_orth(av, xgrid, opt);  
  real(DP), intent(OUT) :: av(:,:);     ! av - array of vectors  
  real(DP), intent(IN) :: xgrid(:);      
  type(LGRR_OPT_T), intent(IN) :: opt;
  integer :: i, mini, maxi;
  real(DP), pointer :: w(:); ! work space
  call dbg_str("-->ho3d_orth()", LGRR_TRACE);
  w => mem_vec_from(xgrid); 
  w = xgrid * opt%lambda;   
  call lgrr_orth_u1(av, w**2, opt%alpha);  ! calling for x=(lambda*r)^2
  call mtrx_i2_range(mini, maxi, av); 
  w = SQRT(w * (opt%lambda * D2));
  do i = mini, maxi  
    av(:, i) = av(:, i) * w; 
  end do 
  call mem_del_vec(w);
end subroutine


function lgrr_lcr3(lcr, maxL, maxNL, lmbdL) result(res)
  type(ARR_VEC_PTR_T), pointer :: res(:);
  type(GRID_LCR_T), intent(IN) :: lcr;
  integer, intent(IN) :: maxL, maxNL(:);
  real(DP), intent(IN) :: lmbdL(:);
  type(LGRR_OPT_T) :: opt;
  real(DP), pointer :: pav(:,:); !temp pointer
  integer :: L;
  character(LEN=64) :: name = "lgrr_lcr3";
  call dbg_in(name);
  call dbg_int("maxL=", maxL);
  call veci_dbg("maxNL=", maxNL);
  call vec_dbg("lmbdL=", lmbdL);
  
  res => mem_veca_ptrs(maxL+1);
  do L = 0, maxL
    opt%lambda = lmbdL(L+1);
    opt%alpha = real(2 * L + 2, DP);
    opt%maxn = maxNL(L+1);
    pav => mem_veca(lcr%xv, opt%maxn+1);
    res(L+1)%av => pav;
    call lgrr_lcr(pav, lcr, opt);
  enddo
  return
end function

subroutine lgrr_lcr(av, lcr, opt);  
  real(DP), intent(OUT) :: av(:,:);     ! av - array of vectors  
  type(GRID_LCR_T), intent(IN) :: lcr;
  type(LGRR_OPT_T), intent(IN) :: opt;
  integer :: i, mini, maxi;
  character(len=1024) :: mssg;
  character(LEN=64) :: name = "lgrr_lcr";
  call dbg_in(name);
  call lgrr_opt_to_str(mssg, opt); 
  call dbg_str2("opt=", TRIM(mssg));
  call lgrr_orth(av, lcr%rv, opt);  !NOTE!!! calculated on rv (not xv)
  call mtrx_i2_range(mini, maxi, av);
  do i = mini, maxi  
    call grid_lcr_wf(av(:, i), lcr);
!    av(:, i) = av(:, i) * lcr%sqcr1; ! 1/sqrt(c+r)
  end do 
end subroutine


subroutine lgrr_lcr_L(av, lcr, opt);  
  real(DP), intent(OUT) :: av(:,:);     ! av - array of vectors  
  type(GRID_LCR_T), intent(IN) :: lcr;
  type(LGRR_OPT_T), intent(INOUT) :: opt;!// L - in; alpha - out
  opt%alpha = real(2 * opt%L + 2, DP);
  call lgrr_lcr(av, lcr, opt);
end subroutine


SUBROUTINE lgrr_opt_to_str(res, opt);    
  character(len=*), intent(OUT) :: res;
  type(LGRR_OPT_T), intent(IN) :: opt;
  write(res, *) "(lambda=", opt%lambda, ", alpha=", opt%alpha &
  & , ", maxn=", opt%maxn, ")";
end subroutine

subroutine ho3d_lcr(av, lcr, opt);  
  real(DP), intent(OUT) :: av(:,:);     ! av - array of vectors  
  type(GRID_LCR_T), intent(IN) :: lcr;
  type(LGRR_OPT_T), intent(IN) :: opt;
  integer :: i, mini, maxi;
  call dbg_str("-->ho3d_lcr()", LGRR_TRACE);   
  call ho3d_orth(av, lcr%rv, opt);  !NOTE!!! calculated on rv (not xv)
  call mtrx_i2_range(mini, maxi, av);
  do i = mini, maxi  
    av(:, i) = av(:, i) * lcr%sqcr1; ! 1/sqrt(c+r)
  end do 
end subroutine

! n=0; NOTE: this is for lambda=1;
real(DP) function lgrr_orth_norm(n, alpha) result(res);   
  real(DP), intent(IN) :: alpha; 
  integer, intent(IN) :: n;
  real(DP) :: dn;
  dn = real(n, DP); 
  res = DLGAMA(dn + D1) - DLGAMA(alpha + dn + D1);
  res = SQRT(EXP(res));
  return;
end function

!!// n=0; NOTE: this is for lambda=1; use lambda*xv for others
SUBROUTINE lgrr_orth_n0(res, xv, alpha);    
  real(DP), intent(OUT) :: res(:);      ! vector-x, vector-y
  real(DP), intent(IN)  :: xv(:), alpha;
  real(DP) :: x, a;     
  integer :: i, mini, maxi;
!  character(LEN=64) :: name = "lgrr_orth_n0";
!  call dbg_in(name);
  call vec_init_res(res, xv);
  call vec_range(mini, maxi, xv); 
!!//      double x = lambda * r;
!!//      if (x == 0)
!!//        return 0;
!!//      return Math.exp(-0.5 * x + 0.5 * alpha * Math.log(x));
  do i = mini, maxi
    x = xv(i);
    if (dbg_is_zero(x) .OR. x < D0) then
      res(i) = D0;
    ELSE
      a = alpha * LOG(x) - x;
      a = D05 * a;
      res(i) = EXP(a);
    end if    
  end do
end SUBROUTINE

end module
